import numpy as np
import pandas as pd
import matplotlib.pyplot as plt

class K_Means:
    def __init__(self, K, X):
        # K簇
        self.C = None

        self.K = K
        self.X = X

        #从X中随机选择K个样本作为初始均值向量{u0 ... uk-1}
        index = np.random.randint(X.shape[0], size = K)
        self.U = X[index]
        self.initial_U = self.U.copy()

    def fit(self):
        round = 0
        while True:
            c = [[] for _ in range(self.K)]

            for x in X:
                d = np.array([np.linalg.norm(x-u) for u in self.U])   #计算距离
                c_x = np.argmin(d)  #距离x最近的簇
                c[c_x].append(x.tolist())   #归入

            # 计算均值向量
            u = []
            for c_i in c:
                u.append(np.mean(c_i, axis=0))
            u = np.array(u)

            print("round:{}, u={}".format(round, u))
            # 更新均值向量
            if (u == self.U).all():
                self.C = c
                break
            else:
                self.U = u

            round += 1
    
            

df = pd.read_csv("./data/watermelon-4.0.csv")
X = df[["density", "sugar"]].to_numpy()

# kmeans = K_Means(2, X)
# kmeans.fit()

# u0 = kmeans.initial_U
# result = kmeans.C
# u = kmeans.U
# print(u)
# c1 = np.array(result[0])
# c2 = np.array(result[1])

# plt.scatter(c1[:,0], c1[:,1], c='red')
# plt.scatter(c2[:,0], c2[:,1], c='blue')
# plt.scatter(u[:,0], u[:,1], c='black', marker='+', s=55,label = "mean")
# plt.scatter(u0[:,0], u0[:,1], c='black', marker='o', s=55,label = "initial_U")
# plt.legend()
# plt.show()



means3 = K_Means(3, X)
means3.fit()

u0 = means3.initial_U
result = means3.C
u = means3.U
print(u)
c1 = np.array(result[0])
c2 = np.array(result[1])
c3 = np.array(result[2])

plt.scatter(c1[:,0], c1[:,1], c='red')
plt.scatter(c2[:,0], c2[:,1], c='blue')
plt.scatter(c3[:,0], c3[:,1], c='green')

plt.scatter(u[:,0], u[:,1], c='black', marker='+', s=55,label = "mean")
plt.scatter(u0[:,0], u0[:,1], c='black', marker='o', s=55,label = "initial_U")
plt.legend()
plt.show()
